#include "sr2.hpp"

#include <vector>

#include "range_coder.hpp"
#include "arith_coder.hpp"
#include "bit_model.hpp"
#include "byte_model.hpp"
#include "qsmodel.hpp"

// An Encoder does arithmetic encoding in n contexts.  Methods:
// Encoder(f, n) creates encoder for compression to archive f, which
//     must be open past any header for writing in binary mode.
// code(cxt, y) compresses bit y (0 or 1) to file f
//   modeled in context cxt (0..n-1)
// flush() should be called exactly once after compression is done and
//     before closing f.

// Create an array p of n elements of type T
template <class T> void alloc(T*&p, int n) {
  p=(T*)calloc(n, sizeof(T));
  //if (!p) quit("out of memory");
}

#define FASTMODEL 1
#define FASTCODER 1


struct StateMap {
  const int N;  // Number of contexts
#if FASTMODEL
  //U16 *t;
  BitModelFast t[(1024+64)*3];
  //ByteModelMultiSymbol bm[1024+64];
  QSModel<256> bm[1024+64];
#else
  U32 *t;       // cxt -> prediction in high 25 bits, count in low 7 bits
#endif
  static int dt[128];  // i -> 1K/(i+2)

  StateMap(int n);

  // predict next bit in context cxt
  int p(int cxt) {
    assert(cxt>=0 && cxt<N);
#if FASTMODEL
	return t[cxt].prob;
#else
    return t[cxt]>>20;
#endif
  }

  // update model in context cxt for actual bit y
  void update(int cxt, int y) {
    assert(cxt>=0 && cxt<N);
    assert(y==0 || y==1);
#if FASTMODEL
	t[cxt].update(y);
#else
    int n=t[cxt]&127, p=t[cxt]>>9;  // count, prediction
    if (n<127) ++t[cxt];
    t[cxt]+=((y<<23)-p)*dt[n]&0xffffff80;
#endif
  }
};

int StateMap::dt[128]={0};

StateMap::StateMap(int n): N(n) {
#if FASTMODEL
	//t = new U16[N];
	//for (int i=0; i<N; ++i)
	//	t[i] = 1 << 11;
#else
	t = new U32[N];
	for (int i=0; i<N; ++i)
		t[i]=1<<31;

	if (dt[0]==0)
		for (int i=0; i<128; ++i)
			dt[i]=512/(i+2);
#endif
}




struct Encoder
{

public:
	Encoder(ByteWriter writer, int n);

	void code(int cxt, int y)
	{
		int p = sm.p(cxt);
		sm.update(cxt, y);
		coder.encode(y, p);
	}

	void flush()
	{
		coder.flush();
	}

	const int N;    // number of contexts
	//FILE* archive;  // Compressed data file
	//std::vector<U8>& archive;
	//U32 x1, x2;     // Range, initially [0, 1), scaled by 2^32
	StateMap sm;    // cxt -> p

#if FASTCODER
	Coder1 coder;
#else
	ACoder coder;
#endif
};

struct Decoder
{
	const int N;     // number of contexts
	//U32 x1, x2;      // Range, initially [0, 1), scaled by 2^32
	//U32 x;           // Decompress mode: last 4 input bytes of archive
	StateMap sm;     // cxt -> p
	Decoder1 decoder;

	Decoder(ByteReader reader, int n);

	int code(int cxt)
	{
		int p = sm.p(cxt);
		int y = decoder.decode(p);
		sm.update(cxt, y);
		return y;
	}
};

Encoder::Encoder(ByteWriter writer, int n)
: N(n), coder(writer), sm(n)
{
}

Decoder::Decoder(ByteReader reader, int n)
: N(n), decoder(reader), sm(n)
{
}

#if 0
// Encode a byte c in context cxt to encoder e
void encode(Encoder& e, int cxt, int c) {

  // code high 4 bits in contexts cxt+1..15
  int b=(c>>4)+16;
  e.code(cxt+1     , b>>3&1);
  e.code(cxt+(b>>3), b>>2&1);
  e.code(cxt+(b>>2), b>>1&1);
  e.code(cxt+(b>>1), b   &1);

  // code low 4 bits in one of 16 blocks of 15 cxts (to reduce cache misses)
  cxt+=15*(b-15);
  b=c&15|16;
  e.code(cxt+1     , b>>3&1);
  e.code(cxt+(b>>3), b>>2&1);
  e.code(cxt+(b>>2), b>>1&1);
  e.code(cxt+(b>>1), b   &1);
}
#endif

// Each byte can potentially be expanded to
// 2*12 + 16 = 40 bit.
// This means a double can use up 40 byte.

// Compress from in to out.  out should be positioned past the header.
void sr2_compress(ByteReader in, ByteWriter out)
{
  std::auto_ptr<Encoder> ep(new Encoder(out, (1024+64)*258));
  Encoder& e = *ep;
  int c1=0;  // previous byte
  U32 h=0;   // hash of last 4 bytes
  U32 *t4;  // context -> last 3 bytes in bits 0..23, count in 24..29
  //alloc(t4, 0x40000);
  alloc(t4, 0x100000);
  U32 pos = 0;
  while (1) {
	if(pos == 1501)
		pos = pos;
    ++pos;
    const U32 r=t4[h];  // last byte count, last 3 bytes in this context
    int cxt;  // context
    if (r>=0x4000000) cxt=1024+(r>>24);
    else cxt=c1|r>>16&0x3f00;
    //cxt*=258;

	int suffix_cxt = cxt * 3;

    if (in.empty()) {  // Mark EOF by coding first match as a literal
      e.code(suffix_cxt, 1);
      e.code(suffix_cxt+1, 0);

	  //cxt = cxt*255 + (1024+64)*3;
      //encode(e, cxt, r&0xff);
	  e.sm.bm[cxt].encode(e.coder, r&0xff);
      break;
    }

	int c = in.get();

    int comp3=c*0x10101^r;  // bytes 0, 1, 2 are 0 if equal
    if (!(comp3&0xff)) {  // match first?
      e.code(suffix_cxt, 0);
      if (r<0x3f000000) t4[h]+=0x1000000;  // increment count
    }
    else if (!(comp3&0xff00)) {  // match second?
      e.code(suffix_cxt, 1);
      e.code(suffix_cxt+1, 1);
      e.code(suffix_cxt+2, 0);
      t4[h]=r&0xff0000|r<<8&0xff00|c|0x1000000;
    }
    else if (!(comp3&0xff0000)) {  // match third?
      e.code(suffix_cxt, 1);
      e.code(suffix_cxt+1, 1);
      e.code(suffix_cxt+2, 1);
      t4[h]=r<<8&0xffff00|c|0x1000000;
    }
    else {  // literal?
	  //printf("literal: %c (%d)\n", (char)c, c);
      e.code(suffix_cxt, 1);
      e.code(suffix_cxt+1, 0);

	  //cxt = cxt*255 + (1024+64)*3;
      //encode(e, cxt, c);
	  e.sm.bm[cxt].encode(e.coder, c);
      t4[h]=r<<8&0xffff00|c;
    }
    c1=c;
    h=h*(5<<5)+c+1&0xfffff;
	//h=(h*160+c)&0x3ffff;
  }
  e.flush();
}

// Decompress from in to out.  in should be positioned past the header.
void sr2_decompress(ByteReader in, ByteWriter out)
{
	std::auto_ptr<Decoder> ep(new Decoder(in, (1024+64)*258));
	Decoder& e = *ep;
	int c1=0;  // previous byte
	U32 h=0;   // hash of last 4 bytes
	U32 *t4;   // context -> last 3 bytes in bits 0..23, count in 24..29
	alloc(t4, 0x100000);

	int pos = 0;

	while (1)
	{
		if(pos == 1501)
			pos = pos;
		++pos;
		const U32 r=t4[h];  // last byte count, last 3 bytes in this context
		int cxt;  // context
		if (r>=0x4000000) cxt=1024+(r>>24);
		else cxt=c1|r>>16&0x3f00;
		
		int suffix_cxt = cxt * 3;

		// Decompress: 0=p[1], 110=p[2], 111=p[3], 10xxxxxxxx=literal.
		// EOF is marked by p[1] coded as a literal.
		if (e.code(suffix_cxt)) {
			if (e.code(suffix_cxt+1)) {
				if (e.code(suffix_cxt+2)) {  // match third?
					c1=r>>16&0xff;
					t4[h]=r<<8&0xffff00|c1|0x1000000;
				}
				else {  // match second?
					c1=r>>8&0xff;
					t4[h]=r&0xff0000|r<<8&0xff00|c1|0x1000000;
				}
			}
			else {  // literal?
				//c1=decode(e, cxt+2);
				c1 = e.sm.bm[cxt].decode(e.decoder);
				if (c1==int(r&0xff)) break;  // EOF?
				t4[h]=r<<8&0xffff00|c1;
			}
		}
		else {  // match first?
			c1=r&0xff;
			if (r<0x3f000000) t4[h]+=0x1000000;  // increment count
		}
		//putc(c1, out);
		out.put(c1);
		h=h*(5<<5)+c1+1&0xfffff;
	}

	out.flush();
}